package net.dybai.glacier;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.crypto.*;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.*;
import java.security.*;
import java.security.spec.AlgorithmParameterSpec;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.KeySpec;
import java.util.Arrays;

/**
 * AES 加密解密工具类
 * @author Added by dybai 2021-02-20.
 */
public class AESUtils {
    private static final Logger LOG = LoggerFactory.getLogger(AESUtils.class);

    // 密钥生成器迭代次数
    private static final int ITERATION_COUNT = 1000;
    // 密钥长度，AES 支持 128、192 和 256 bit 的密钥
    private static final int KEY_LENGTH = 256;
    // IV 长度，AES IV 必须是 128 bit
    private static final int IV_LENGTH = 128;
    // 盐值长度，Byte
    private static final int SALT_LENGTH = 8;

    // 各种标准名称详见 Java Cryptography Architecture Reference Guide 的附录 A：
    // https://docs.oracle.com/javase/8/docs/technotes/guides/security/crypto/CryptoSpec.html

    // 密钥加密算法
    private static final String KEY_ALGORITHM = "PBKDF2WithHmacSHA256";
    // 数据加密算法
    private static final String DATA_ALGORITHM = "AES";
    // 操作模式
    private static final String OPERATION_MODULES = "CFB";
    // 填充方案
    private static final String PADDING_SCHEME = "PKCS5Padding";
    // 数据加密算法转换名称
    private static final String DATA_ALG_TRANSFORMATION = DATA_ALGORITHM + "/"
            + OPERATION_MODULES + "/" + PADDING_SCHEME;

    // 文件魔数
    private static final byte[] MAGIC = new byte[] {
            //'S', 'a', 'l', 't', 'e', 'd', '_', '_',
            0x53, 0x61, 0x6C, 0x74, 0x65, 0x64, 0x5F, 0x5F
    };
    // 加密解密文件缓冲大小
    private static final int BUFSIZE = 1024 * 8; // 8KB

    private enum Operation {
        ENCRYPT, DECRYPT
    }

    /**
     * 生成盐值
     * @return
     */
    private static byte[] generateSalt() {
        byte[] salt = new byte[SALT_LENGTH];
        SecureRandom random = new SecureRandom();
        random.nextBytes(salt);
        return salt;
    }

    /*
     * 生成 AES 密钥。<br/>
     * 此函数采用 MD5 算法生成密钥，因此只能满足采用 128bit 密钥的加密方式使用。
     */
//    private static SecretKeySpec getAESKey(char[] password, byte[] salt)
//            throws NoSuchAlgorithmException {
//        byte[] pwd_byte = Utils.convertCharArrayToByteArray(password);
//        byte[] data = new byte[pwd_byte.length + salt.length];
//        byte[] key = null;
//
//        System.arraycopy(pwd_byte, 0, data, 0, pwd_byte.length);
//        System.arraycopy(salt, 0, data, pwd_byte.length, salt.length);
//
//        MessageDigest md = MD5Utils.beginDigest();
//        MD5Utils.doDigest(md, data, 0, data.length);
//        key = MD5Utils.endDigest(md);
//        return new SecretKeySpec(key, DATA_ALGORITHM);
//    }

    /**
     * 生成 AES 密钥<br/>
     * @param password 用户设置的密码
     * @param salt 盐值
     * @return AES 密钥
     * @throws InvalidKeySpecException
     * @throws NoSuchAlgorithmException
     */
    private static SecretKeySpec getAESKey(char[] password, byte[] salt)
            throws InvalidKeySpecException, NoSuchAlgorithmException {
        // 生成 PBE 密钥
        KeySpec keySpec = new PBEKeySpec(password, salt, ITERATION_COUNT, KEY_LENGTH);
        // 通过工厂获取密钥
        SecretKeyFactory keyFactory = SecretKeyFactory.getInstance(KEY_ALGORITHM);
        byte[] key = keyFactory.generateSecret(keySpec).getEncoded();
        return new SecretKeySpec(key, DATA_ALGORITHM);
    }

//    private static byte[] getIV(SecretKeySpec aesKey, char[] password
//            , byte[] salt) throws NoSuchAlgorithmException {
//        byte[] key = aesKey.getEncoded();
//        byte[] pwd_byte = Utils.convertCharArrayToByteArray(password);
//        byte[] data = new byte[key.length + pwd_byte.length + salt.length];
//
//        System.arraycopy(key, 0, data, 0, key.length);
//        System.arraycopy(pwd_byte, 0, data, key.length, pwd_byte.length);
//        System.arraycopy(salt, 0, data, key.length + pwd_byte.length
//                , salt.length);
//
//        MessageDigest md = MD5Utils.beginDigest();
//        MD5Utils.doDigest(md, data, 0, data.length);
//        return MD5Utils.endDigest(md);
//    }

    /**
     * 生成 IV<br/>
     * 由于 IV 总是 128bit，因此采用 PBEKeySpec 来生成还是采用 MD5 的方式来生成均可。
     * @param aesKey AES 密钥
     * @param password 用户设置的密码
     * @param salt 盐值
     * @return CFB、CBC 加密方式必须使用此值参与运算。
     * @throws NoSuchAlgorithmException
     * @throws InvalidKeySpecException
     */
    private static byte[] getIV(SecretKeySpec aesKey, char[] password
            , byte[] salt) throws NoSuchAlgorithmException, InvalidKeySpecException {
        byte[] key = aesKey.getEncoded();
        byte[] data = new byte[key.length + salt.length];

        System.arraycopy(key, 0, data, 0, key.length);
        System.arraycopy(salt, 0, data, key.length, salt.length);

        // 生成 PBE 密钥
        KeySpec keySpec = new PBEKeySpec(password, data, ITERATION_COUNT, IV_LENGTH);
        // 通过工厂获取密钥
        SecretKeyFactory keyFactory = SecretKeyFactory.getInstance(KEY_ALGORITHM);
        return keyFactory.generateSecret(keySpec).getEncoded();
    }

    /**
     * 初始化加密解密器
     * @param op 操作方向
     * @param password 用户设置的密码
     * @param salt 盐值
     * @return 加密解密器
     * @throws InvalidKeySpecException
     * @throws NoSuchAlgorithmException
     * @throws NoSuchPaddingException
     * @throws InvalidKeyException
     * @throws InvalidAlgorithmParameterException
     */
    private static Cipher initCipher(Operation op, char[] password, byte[] salt, boolean isShowSecret)
            throws InvalidKeySpecException, NoSuchAlgorithmException
            , NoSuchPaddingException, InvalidKeyException
            , InvalidAlgorithmParameterException {
        int mode = -1;
        // 根据用户密码和盐值获取 AES KEY
        SecretKeySpec aesKey = getAESKey(password, salt);
        // 获取 IV
        byte[] iv = getIV(aesKey, password, salt);
        AlgorithmParameterSpec ivParameterSpec = new IvParameterSpec(iv);

        if (isShowSecret) {
            LOG.debug("salt = " + HexUtils.toHexString(salt));
            LOG.debug("key = " + HexUtils.toHexString(aesKey.getEncoded()));
            LOG.debug("iv = " + HexUtils.toHexString(iv));
        }

        // 创建加密解密器
        Cipher cipher = Cipher.getInstance(DATA_ALG_TRANSFORMATION);

        switch (op) {
            case ENCRYPT:
                mode = Cipher.ENCRYPT_MODE;
                break;
            case DECRYPT:
                mode = Cipher.DECRYPT_MODE;
                break;
            default:
                throw new InvalidParameterException("Invalid operation mode");
        }
        // 初始化加密解密器
        cipher.init(mode, aesKey, ivParameterSpec);
        return cipher;
    }

    /**
     * 开始加密，后面需要调用 {@see doOperation(Cipher cipher, byte[] data<br/>
     * , int offset, int len)} 函数进行加密运算，并在所有数据处理完成后调用<br/>
     * {@see endOperation(Cipher cipher)} 函数结束加密运算。
     * @param password 用户密码
     * @param salt 盐值
     * @return 加密器
     * @throws InvalidKeySpecException
     * @throws NoSuchAlgorithmException
     * @throws NoSuchPaddingException
     * @throws InvalidKeyException
     * @throws InvalidAlgorithmParameterException
     */
    public static Cipher beginEncrypt(char[] password, byte[] salt, boolean isShowSecret)
            throws InvalidKeySpecException, NoSuchAlgorithmException
            , NoSuchPaddingException, InvalidKeyException
            , InvalidAlgorithmParameterException {
        return initCipher(Operation.ENCRYPT, password, salt, isShowSecret);
    }

    /**
     * 开始解密，后面需要调用 {@see doOperation(Cipher cipher, byte[] data < br / >
     **, int offset, int len)} 函数进行解密运算，并在所有数据处理完成后调用<br/>
     *      * {@see endOperation(Cipher cipher)} 函数结束解密运算。
     * @param password 用户密码
     * @param salt 盐值
     * @return 解密器
     * @throws InvalidKeySpecException
     * @throws NoSuchAlgorithmException
     * @throws NoSuchPaddingException
     * @throws InvalidKeyException
     * @throws InvalidAlgorithmParameterException
     */
    public static Cipher beginDecrypt(char[] password, byte[] salt, boolean isShowSecret)
            throws InvalidKeySpecException, NoSuchAlgorithmException
            , NoSuchPaddingException, InvalidKeyException
            , InvalidAlgorithmParameterException {
        return initCipher(Operation.DECRYPT, password, salt, isShowSecret);
    }

    /**
     * 加密或解密运算，具体是执行加密运算还是解密运算，由 cipher 参数决定。
     * @param cipher 加密解密器
     * @param data 如果是加密操作，此参数应为明文；如果是解密操作，此参数应为密文；
     * @param offset data 参数参与运算的起始位置；
     * @param len data 参数参与运算的数据长度；
     * @return 运算结果。如果是加密操作，此值为密文；如果是解密操作，此值为明文；
     * @throws BadPaddingException
     * @throws IllegalBlockSizeException
     */
    public static byte[] doOperation(Cipher cipher, byte[] data
            , int offset, int len) throws BadPaddingException
            , IllegalBlockSizeException {
        return cipher.update(data, offset, len);
    }

    /**
     * 结束加密或解密运算
     * @param cipher
     * @return
     * @throws BadPaddingException
     * @throws IllegalBlockSizeException
     */
    public static byte[] endOperation(Cipher cipher) throws BadPaddingException
            , IllegalBlockSizeException {
        return cipher.doFinal();
    }

    /**
     * 文件加密
     * @param in 明文文件
     * @param out 密文文件
     * @param password 密码
     * @throws IOException
     * @throws NoSuchAlgorithmException
     * @throws InvalidKeyException
     * @throws InvalidKeySpecException
     * @throws NoSuchPaddingException
     * @throws BadPaddingException
     * @throws IllegalBlockSizeException
     * @throws InvalidAlgorithmParameterException
     */
    public static void fileEncrypt(InputStream in, OutputStream out, String password, boolean isShowSecret)
            throws IOException, NoSuchAlgorithmException
            , InvalidKeyException, InvalidKeySpecException
            , NoSuchPaddingException, BadPaddingException
            , IllegalBlockSizeException, InvalidAlgorithmParameterException {
        int readLen = -1;
        byte[] plainBuf = new byte[BUFSIZE];
        byte[] encryptBuf = null;
        byte[] salt = generateSalt();

        Cipher cipher = beginEncrypt(password.toCharArray(), salt, isShowSecret);

        // 加密后的文件1-8字节为魔数，9-16字节为盐值
        out.write(MAGIC); // 写入魔数
        out.write(salt); // 写入盐值

        while ((readLen = in.read(plainBuf)) >= 0) {
            encryptBuf = doOperation(cipher, plainBuf, 0, readLen);
            out.write(encryptBuf);
        }
        encryptBuf = endOperation(cipher);
        out.write(encryptBuf);
        out.flush();
    }

    /**
     * 文件加密
     * @param inFile 明文文件
     * @param outFile 密文文件
     * @param password 密码
     * @throws IOException
     * @throws NoSuchAlgorithmException
     * @throws InvalidKeyException
     * @throws InvalidKeySpecException
     * @throws NoSuchPaddingException
     * @throws BadPaddingException
     * @throws IllegalBlockSizeException
     * @throws InvalidAlgorithmParameterException
     */
    public static void fileEncrypt(File inFile, File outFile, String password, boolean isShowSecret)
            throws IOException, NoSuchAlgorithmException
            , InvalidKeyException, InvalidKeySpecException
            , NoSuchPaddingException, BadPaddingException
            , IllegalBlockSizeException, InvalidAlgorithmParameterException {
        try (InputStream in = new FileInputStream(inFile);
             OutputStream out = new FileOutputStream(outFile)) {
            fileEncrypt(in, out, password, isShowSecret);
        }
    }

    /**
     * 检查加密文件的魔数
     * @param in 加密文件输入流
     * @return true 表示校验通过；false 表示校验失败。
     * @throws IOException
     */
    private static boolean checkFileMagic(InputStream in) throws IOException {
        int readLen = -1;
        byte[] buf = new byte[SALT_LENGTH];

        readLen = in.read(buf);
        if (SALT_LENGTH != readLen) {
            throw new IOException("[MAGIC] Invalid file length.");
        }

        return Arrays.equals(MAGIC, buf);
    }

    /**
     * 文件解密
     * @param inFile 密文文件
     * @param outFile 明文文件
     * @param password 密码
     * @throws IOException
     * @throws NoSuchAlgorithmException
     * @throws InvalidKeyException
     * @throws InvalidKeySpecException
     * @throws NoSuchPaddingException
     * @throws BadPaddingException
     * @throws IllegalBlockSizeException
     * @throws InvalidAlgorithmParameterException
     */
    public static void fileDecrypt(File inFile, File outFile, String password, boolean isShowSecret)
            throws IOException, NoSuchAlgorithmException, InvalidKeyException
            , InvalidKeySpecException, NoSuchPaddingException
            , BadPaddingException, IllegalBlockSizeException
            , InvalidAlgorithmParameterException {
        int readLen = -1;
        byte[] encryptBuf = new byte[BUFSIZE];
        byte[] plainBuf = null;
        byte[] salt = new byte[SALT_LENGTH];

        try (InputStream in = new FileInputStream(inFile);
             OutputStream out = new FileOutputStream(outFile)) {

            // 校验魔数
            if (!checkFileMagic(in)) {
                throw new IOException("[MAGIC] Invalid encrypt file");
            }

            // 读取盐值
            readLen = in.read(salt);
            if (SALT_LENGTH != readLen) {
                throw new IOException("[SALT] Invalid file length.");
            }

            Cipher cipher = beginDecrypt(password.toCharArray(), salt, isShowSecret);
            while ((readLen = in.read(encryptBuf)) >= 0) {
                plainBuf = doOperation(cipher, encryptBuf, 0, readLen);
                out.write(plainBuf);
            }
            plainBuf = endOperation(cipher);
            out.write(plainBuf);
            out.flush();
        }
    }

}
